#!/bin/bash
# Copyright (C) 2024-2025 KoudaiPC Computer Technologies Co., Ltd.
# script name: yandi-populate-root.sh
# author: david921518@qq.com

################
# global vars
GLB_WORK_DIR=
GLB_DONE_FILE=
GLB_FIRSTBOOT_ROOT=
GLB_ROOT_DEV=
GLB_ROOT_PART=
GLB_ADMIN_SSH_PORT=
GLB_ADMIN_PASSWORD=
GLB_ADMIN_USERNAME=
GLB_LC_LANGUAGE=
GLB_LC_CHARSET=
################

user=$(whoami)

# 定义日志格式-前缀
log_prefix="[populate-root]"
# 定义日志目录和文件名
log_file_dir="/var/log/kpc"
log_file_path="$log_file_dir/firstboot.log"

log_kmsg_path="/dev/kmsg"
log_con_path="/dev/tty1"

function log_init()
{
  # 文件不存在则创建
  if [ ! -d "$log_file_dir" ]; then
    mkdir -p $log_file_dir
    chmod -R 755 $log_file_dir
  fi
  if [ ! -f "$log_file_path" ]; then
    touch $log_file_path
  fi
}

# 定义日志打印格式
function log()
{
  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")
  echo "$curtime $log_prefix $@"
  if [ "$user" == "root" ]; then
    echo "$curtime $log_prefix $@" >>$log_file_path
    #sleep 0.1
    echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    sleep 0.1
    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #sleep 0.1
  fi
}

function log_from_file()
{
  tmpfile=$@
  if [ ! -f "$tmpfile" ]; then
    log "no file [$tmpfile] to log"
    return 1
  fi

  local n=0
  local -a lists
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$tmpfile

  # 定义日志时间格式
  local curtime=$(date "+%Y-%m-%d %H:%M:%S")

  #echo "$curtime $log_prefix $@"
  local i=0
  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    echo "$curtime $log_prefix ${lists[i]}"
  done

  if [ "$user" == "root" ]; then
    #echo "$curtime $log_prefix $@" >>$log_file_path
    for i in $( seq 0 $((${#lists[*]} - 1)) )
    do
      echo "$curtime $log_prefix ${lists[i]}" >>$log_file_path
    done

    #sleep 0.1

    #echo "$log_prefix [kmsg] $@" >$log_kmsg_path
    for i in $( seq 0 $((${#lists[*]} - 1)) )
    do
      echo "$curtime $log_prefix ${lists[i]}" >$log_kmsg_path
      [ $((i%10)) -ne 9 ] || sleep 0.1
    done

    #sleep 0.1

    #echo "$curtime $log_prefix [con] $@" >$log_con_path
    #for i in $( seq 0 $((${#lists[*]} - 1)) )
    #do
    #  echo "$curtime $log_prefix ${lists[i]}" >$log_con_path
    #done

    #sleep 0.1
  fi
}

# loop forever to reboot
function force_reboot()
{
  log "force reboot now!"
  while true; do reboot; sleep 9; done
}

# fix all kinds of possible broken issues
function fix_possible_broken()
{
  log "apt-get --fix-broken install -y"
  apt-get --fix-broken install -y

  log "dpkg -C"
  dpkg -C
  log "dpkg --configure -a"
  dpkg --configure -a
}

function get_work_dir()
{
  ################
  # /kpcos/distro/yandi/releases/0.5/ports/arm64/guest/yandi-populate-root.sh
  # /kpcos/distro/yandi/releases/0.5/ports/arm64/root/etc/zzz-firstboot
  ################

  local real_path=$(realpath $0)
  GLB_WORK_DIR=$(dirname "$real_path")
  log "WORK DIR is [$GLB_WORK_DIR]"

  local done_file="${real_path}_done"
  log "done_file=[$done_file]"
  real_done_file=$(realpath "$done_file")
  log "real_done_file=[$real_done_file]"

  GLB_DONE_FILE="$real_done_file"
  log "GLB_DONE_FILE=[$GLB_DONE_FILE]"

  local firstboot_root="$GLB_WORK_DIR/../root"
  log "firstboot_root=[$firstboot_root]"
  real_firstboot_root=$(realpath "$firstboot_root")
  log "real_firstboot_root=[$real_firstboot_root]"

  if [ ! -d $real_firstboot_root ]; then
    log "No firstboot root dir [$real_firstboot_root]."
    exit 1
  fi

  GLB_FIRSTBOOT_ROOT="$real_firstboot_root"
  log "GLB_FIRSTBOOT_ROOT=[$GLB_FIRSTBOOT_ROOT]"
}

# save installed packages list
function save_pkg_list()
{
  if [ "$1" == "" ]; then
    log "save_pkg_list() filename is empty, do nothing!"
    return
  fi
  local list_file="$PWD"/$1
  log "save bootstrap packages list to $list_file"
  [ -f $list_file ] || dpkg -l >$list_file
}

# get ROOT device from mount info
function get_root_dev()
{
  local line=$(mount | grep " on / type ")
  if [ x"$line" != x ]; then
    root_part=$(echo $line | awk '{print $1}')
    case "$root_part" in
      "/dev/sd"* )
        GLB_ROOT_DEV=${root_part:0:8}
        GLB_ROOT_PART=$root_part
        ;;
      "/dev/vd"* )
        GLB_ROOT_DEV=${root_part:0:8}
        GLB_ROOT_PART=$root_part
        ;;
      * )
        ;;
    esac
    log "GLB_ROOT_DEV=[$GLB_ROOT_DEV] GLB_ROOT_PART=[$GLB_ROOT_PART]"
  fi
}

# start kpc-heartbeat
function start_heartbeat()
{
  heartbeat_conf=/etc/kpc/heartbeat.conf
  heartbeat_file=/usr/sbin/kpc-heartbeat.sh
  heartbeat_lock=/tmp/kpc-hearbeat.lock
  if [ -x $heartbeat_file ]; then
    log "start kpc-heartbeat"
    (flock -n $heartbeat_lock -c "echo 9 >$heartbeat_conf ; ($heartbeat_file)")&
  fi
}

# run bootstrap configurating
function run_bootstrap_config()
{
  local admin_ssh_port=""
  local admin_password=""
  local admin_username=""
  local lc_language=""
  local lc_charset=""
  local wifi_backend=""
  local wifi_dev=""
  local wifi_addr=""
  local wifi_mode=""
  local wifi_net=""
  local wifi_pass=""

  local conf_file=/etc/kpc/bootstrap.conf
  if [ -f $conf_file ]; then
    log "start to configurate bootstrap environment"

    local n=0
    local -a lists
    while read line;
    do
      lists[$n]="$line"
      ((n++))
    done <$conf_file

    local i=0
    for i in $( seq 0 $((${#lists[*]} - 1)) )
    do
      #log "lists[$i]=[${lists[i]}]"
      local name=$(echo ${lists[i]} | awk '{print $1}')
      local val=$(echo ${lists[i]} | awk '{print $2}')
      #log "name=[$name] val=[$val]"
      case "$name" in
        "#"* )
          #log "it's comment"
          ;;
        "" )
          #log "it's empty line"
          ;;
        "admin.ssh_port" )
          admin_ssh_port="$val"
          if [ $admin_ssh_port -gt 65535 ]; then
            log "admin.ssh_port=[$admin_ssh_port] is invalid"
            continue;
          fi
          if [ $admin_ssh_port -lt 22 ]; then
            log "admin.ssh_port=[$admin_ssh_port] is invalid"
            continue;
          fi
          log "admin SSH port set to [$admin_ssh_port]"
          GLB_ADMIN_SSH_PORT=$admin_ssh_port
          ;;
        "admin.password" )
          admin_password="$val"
          log "admin pasword set to [$admin_password]"
          GLB_ADMIN_PASSWORD=$admin_password
          ;;
        "admin.username" )
          admin_username="$val"
          log "admin username set to [$admin_username]"
          GLB_ADMIN_USERNAME=$admin_username
          ;;
        "locale.language" )
          lc_language="$val"
          log "LANGUAGE set to [$lc_language]"
          GLB_LC_LANGUAGE=$lc_language
          ;;
        "locale.charset" )
          lc_charset="$val"
          log "CHARSET set to [$lc_charset]"
          GLB_LC_CHARSET=$lc_charset
          ;;
        "wifi.backend" )
          if [ "$val" == "iwd" ]; then
            log "enable iwd network configuration"
            wifi_backend="$val"
            iwd_main_conf=/etc/iwd/main.conf
            sed "s/#EnableNetworkConfiguration=true/EnableNetworkConfiguration=true/g" -i $iwd_main_conf
            #log_from_file $iwd_main_conf
          fi
          ;;
        "wifi.device_name" )
          wifi_dev="$val"
          log "wifi_dev=[$wifi_dev]"
          ;;
        "wifi.device_address" )
          wifi_addr="$val"
          log "wifi_addr=[$wifi_addr]"
          ;;
        "wifi.mode" )
          wifi_mode="$val"
          log "wifi_mode=[$wifi_mode]"
          ;;
        "wifi.network" )
          wifi_net="$val"
          log "wifi_net=[$wifi_net]"
          ;;
        "wifi.passphrase" )
          wifi_pass="$val"
          log "wifi_pass=[$wifi_pass]"
          ;;
        * )
          log "unknown option=[$name]"
          ;;
      esac
    done

    if [ "$wifi_backend" == "iwd" ]; then
      log "systemctl --now enable iwd"
      systemctl --now enable iwd
      sleep 3
      if [ x"$wifi_dev" == x ]; then
        #iwctl device list
        line=$(iwctl device list | grep "station" )
        log "line=[$line]"
        if [ x"$line" == x ]; then
          log "no station mode wifi device found, exit now"
          exit 1
        fi
        #dummy=$(echo $line | awk '{print $1}')
        dummy=$(echo $line | awk '{print $1}')
        log "dummy=[$dummy]"
        #dev_name=$(echo $line | awk '{print $2}')
        wifi_dev=$(echo $line | awk '{print $2}')
        log "wifi_dev=[$wifi_dev]"
        #dev_addr=$(echo $line | awk '{print $3}')
        wifi_addr=$(echo $line | awk '{print $3}')
        log "wifi_addr=[$wifi_addr]"
        #dev_powered=$(echo $line | awk '{print $4}')
        wifi_powered=$(echo $line | awk '{print $4}')
        log "wifi_powered=[$wifi_powered]"
        #dev_adapter=$(echo $line | awk '{print $5}')
        wifi_adapter=$(echo $line | awk '{print $5}')
        log "wifi_adapter=[$wifi_adapter]"
        #dev_mode=$(echo $line | awk '{print $6}')
        wifi_mode=$(echo $line | awk '{print $6}')
        log "wifi_mode=[$wifi_mode]"

        log "detect wifi_dev=[$wifi_dev]"
        if [ x"$wifi_dev" == x ]; then
          log "no wifi device found, exit now"
          exit 1
        fi
      fi
      if [ x"$wifi_dev" != x ]; then
        line=$(iwctl device list | grep " $wifi_dev " )
        log "line=[$line]"
        if [ x"$line" != x ]; then
          dummy=$(echo $line | awk '{print $1}')
          dev_name=$(echo $line | awk '{print $2}')
          dev_addr=$(echo $line | awk '{print $3}')
          dev_powered=$(echo $line | awk '{print $4}')
          dev_adapter=$(echo $line | awk '{print $5}')
          dev_mode=$(echo $line | awk '{print $6}')
          log "dev_mode=[$dev_mode]"
          if [ "$dev_mode" == "station" ]; then
            #iwctl station wlan0 scan
            tmpfile=$(mktemp)
            iwctl station $wifi_dev scan 2>&1 | tee $tmpfile
            log_from_file $tmpfile
            rm -f $tmpfile
            sleep 3

            #iwctl station wlan0 get-networks
            tmpfile=$(mktemp)
            iwctl station $wifi_dev get-networks 2>&1 | tee $tmpfile
            log_from_file $tmpfile
            rm -f $tmpfile

            #iwctl --passphrase "mysecret" station wlan0 connect "SSID"
            log "iwctl --passphrase $wifi_pass station $wifi_dev connect $wifi_net"
            iwctl --passphrase "$wifi_pass" station $wifi_dev connect "$wifi_net"
            sleep 3

            #iwctl station wlan0 show
            tmpfile=$(mktemp)
            iwctl station $wifi_dev show 2>&1 | tee $tmpfile
            log_from_file $tmpfile
            rm -f $tmpfile
            log "/etc/init.d/iwd restart"
            /etc/init.d/iwd restart
            sleep 1
          fi
        fi
      fi
      #iwctl station wlan0 scan
      #iwctl station wlan0 get-networks
      #iwctl --passphrase "mysecret" station wlan0 connect "SSID"
      #iwctl station wlan0 show
      #/etc/init.d/iwd restart
      #iwctl station wlan0 disconnect
    fi #endof wifi_backend==iwd
  fi
}

# check WAN interface
function check_wan_iface()
{
  local COUNTER=0
  local dev_array=($(ip route | awk '/default/ {print $5}'))
  local iface_dev=${dev_array[0]}
  while [ "$iface_dev"x == x ]; do
    sleep 1
    ((COUNTER++))
    if [ $COUNTER -gt 29 ]; then
      log "Can't get default route device interface, exit now"
      exit 1
    fi
    dev_array=($(ip route | awk '/default/ {print $5}'))
    iface_dev=${dev_array[0]}
  done

  COUNTER=0
  local WAN_IP=$(ip addr show $iface_dev | grep "inet" | head -n 1 | awk '/inet/ {print $2}' | cut -d'/' -f1)
  while [ "${WAN_IP}"x == x ]; do
    sleep 1
    ((COUNTER++))
    if [ $COUNTER -gt 29 ]; then
      log "Can't get WAN IP address, reboot now"
      force_reboot
    fi
    log "get WAN IP address again"
    WAN_IP=$(ip addr show $iface_dev | grep "inet" | head -n 1 | awk '/inet/ {print $2}' | cut -d'/' -f1)
  done
  log "WAN IP=[${WAN_IP}]"
}

function purge_removed_pkgs()
{
  local tmpfile=$(mktemp)
  dpkg -l >$tmpfile
  local n=0
  local -a lists
  while read line;
  do
    lists[$n]="$line"
    ((n++))
  done <$tmpfile
  rm -f $tmpfile

  local i=0
  for i in $( seq 0 $((${#lists[*]} - 1)) )
  do
    local pkg_flag=$(echo "${lists[i]}" | awk '{print $1}')
    local pkg_name=$(echo "${lists[i]}" | awk '{print $2}')
    if [ "${pkg_flag}" == "rc" ]; then
      log "apt-get purge $pkg_name"
      tmpfile=$(mktemp)
      apt-get purge -y $pkg_name 2>&1 | tee $tmpfile
      #log_from_file $tmpfile
      rm -f $tmpfile
    fi
  done
}

# 0. init script
log_init
log "[populate-root] START"
fix_possible_broken
get_work_dir
if [ -f $GLB_DONE_FILE ]; then
  log "exist file $GLB_DONE_FILE , exit now"
  exit 1
fi
get_root_dev
start_heartbeat
run_bootstrap_config
check_wan_iface
save_pkg_list pre_populate_root_pkgs.txt

log "PATH=[$PATH]"

log "fix potential dpkg issues"
dpkg -C
dpkg --configure -a

# 1. banner
if [ -f "$GLB_FIRSTBOOT_ROOT/etc/banner" ]; then
  log "copy $GLB_FIRSTBOOT_ROOT/etc/banner to /etc/banner"
  cp -f "$GLB_FIRSTBOOT_ROOT/etc/banner" /etc/banner
fi

# 2. /etc
if [ -e "$GLB_FIRSTBOOT_ROOT/etc" ]; then
  log "copy $GLB_FIRSTBOOT_ROOT/etc/* to /etc/"
  cp -af "$GLB_FIRSTBOOT_ROOT"/etc/* /etc/
fi

# 3. /lib
if [ -e "$GLB_FIRSTBOOT_ROOT/lib" ]; then
  log "copy $GLB_FIRSTBOOT_ROOT/lib/* to /lib/"
  cp -af "$GLB_FIRSTBOOT_ROOT"/lib/* /lib/
fi

# 4. /usr/sbin
if [ -e "$GLB_FIRSTBOOT_ROOT/usr/sbin" ]; then
  log "copy $GLB_FIRSTBOOT_ROOT/usr/sbin/* to /usr/sbin/"
  cp -af "$GLB_FIRSTBOOT_ROOT"/usr/sbin/* /usr/sbin/
fi

# 5. purge removed pkgs
log "purge removed pkgs"
#if [ -x /usr/sbin/kpc-apt-purge-all-removed-pkgs.sh ]; then
#  /usr/sbin/kpc-apt-purge-all-removed-pkgs.sh
#fi
purge_removed_pkgs

# 6. add service-keeper in cron
sys_crontab=/etc/crontab
result=$(grep "kpc-service-keeper.sh" $sys_crontab)
if [ x"$result" == x ]; then
  echo " * *    * * *   root    test -x /usr/sbin/kpc-service-keeper.sh && /usr/sbin/kpc-service-keeper.sh" >>$sys_crontab
fi

# 7. populate kpc user flatpak remote repos
if [ x"$GLB_ADMIN_USERNAME" != x ]; then
  def_username="$GLB_ADMIN_USERNAME"
else
  def_username="kpc"
fi

log "populate system flatpak remote repos"
flatpak remote-add --system --if-not-exists flathub https://dl.flathub.org/repo/flathub.flatpakrepo
flatpak remote-modify --system flathub --url=https://mirrors.ustc.edu.cn/flathub
flatpak install --system -y flathub io.github.flattool.Warehouse

let COUNTER=0
warehouse_installed=$(flatpak list --system --app --columns=application | grep "io.github.flattool.Warehouse")
while [ "$warehouse_installed" != "io.github.flattool.Warehouse" ]; do
  sleep 1
  let COUNTER=$COUNTER+1
  if [ $COUNTER -gt 9 ]; then
    log "Can't install io.github.flattool.Warehouse, exit now"
    exit 1
  fi
  flatpak remote-add --system --if-not-exists flathub https://dl.flathub.org/repo/flathub.flatpakrepo
  flatpak remote-modify --system flathub --url=https://mirrors.ustc.edu.cn/flathub
  flatpak install --system -y flathub io.github.flattool.Warehouse
  warehouse_installed=$(flatpak list --system --app --columns=application | grep "io.github.flattool.Warehouse")
done

tmpfile=$(mktemp)
chmod 0666 $tmpfile
flatpak remotes --system 2>&1 | tee $tmpfile
log_from_file $tmpfile
rm -f $tmpfile

if [ "yes" == "no" ]; then
log "populate [$def_username] user flatpak remote repos"
su - "$def_username" -c "flatpak remote-add --user --if-not-exists flathub https://dl.flathub.org/repo/flathub.flatpakrepo"
su - "$def_username" -c "flatpak remote-modify flathub --url=https://mirrors.ustc.edu.cn/flathub"
su - "$def_username" -c "flatpak install --user -y flathub io.github.flattool.Warehouse"

let COUNTER=0
warehouse_installed=$(su - "$def_username" -c "flatpak list --user --app --columns=application" | grep "io.github.flattool.Warehouse")
while [ "$warehouse_installed" != "io.github.flattool.Warehouse" ]; do
  sleep 1
  let COUNTER=$COUNTER+1
  if [ $COUNTER -gt 9 ]; then
    log "Can't install io.github.flattool.Warehouse, exit now"
    exit 1
  fi
  su - "$def_username" -c "flatpak remote-add --user --if-not-exists flathub https://dl.flathub.org/repo/flathub.flatpakrepo"
  su - "$def_username" -c "flatpak remote-modify flathub --url=https://mirrors.ustc.edu.cn/flathub"
  su - "$def_username" -c "flatpak install --user -y flathub io.github.flattool.Warehouse"
  warehouse_installed=$(su - "$def_username" -c "flatpak list --user --app --columns=application" | grep "io.github.flattool.Warehouse")
done

tmpfile=$(mktemp)
chmod 0666 $tmpfile
su - "$def_username" -c "flatpak remotes 2>&1 | tee $tmpfile"
log_from_file $tmpfile
rm -f $tmpfile
fi
# end of if [ "yes" == "no" ]; then

# 8. /home/kpc/.config/
# FIXME: default user is kpc, and we have $GLB_FIRSTROOT_ROOT/home/kpc/ directory!!!
if [ -e "$GLB_FIRSTBOOT_ROOT/home/kpc/.config" ]; then
  log "copy $GLB_FIRSTBOOT_ROOT/home/kpc/.config/* to /home/$def_username/.config/"
  if [ ! -e "/home/$def_username"/.config ]; then
    mkdir -p "/home/$def_username"/.config
    chmod 0700 "/home/$def_username"/.config
  fi
  cp -af "$GLB_FIRSTBOOT_ROOT"/home/kpc/.config/* "/home/$def_username"/.config/
  chown "${def_username}:${def_username}" "/home/$def_username"/.config -R
fi

# 9. setup NetworkManager iwd backend
conf_file=/etc/NetworkManager/conf.d/iwd.conf
touch $conf_file
if [ x"$(grep 'wifi.backend=iwd' $conf_file)" == x ]; then
  echo "[device]" >>$conf_file
  echo "wifi.backend=iwd" >>$conf_file
fi

systemctl stop NetworkManager
sleep 0.1
systemctl disable --now wpa_supplicant
sleep 0.1
systemctl restart NetworkManager

# 10. setup dropbear to SSH admin port
if [ x"$GLB_ADMIN_SSH_PORT" == x ]; then
  # use default port 50022
  log "change dropbear port to 50022"
  #sed "s/#DROPBEAR_PORT=22/DROPBEAR_PORT=50022/g" -i /etc/default/dropbear
  sed "s/DROPBEAR_PORT=22/DROPBEAR_PORT=50022/g" -i /etc/default/dropbear
  sed "s/DROPBEAR_PORT=22/DROPBEAR_PORT=50022/g" -i /etc/init.d/dropbear
else
  # use setting port $GLB_ADMIN_SSH_PORT
  log "change dropbear port to $GLB_ADMIN_SSH_PORT"
  #sed "s/#DROPBEAR_PORT=22/DROPBEAR_PORT=$GLB_ADMIN_SSH_PORT/g" -i /etc/default/dropbear
  sed "s/DROPBEAR_PORT=22/DROPBEAR_PORT=$GLB_ADMIN_SSH_PORT/g" -i /etc/default/dropbear
  sed "s/DROPBEAR_PORT=22/DROPBEAR_PORT=$GLB_ADMIN_SSH_PORT/g" -i /etc/init.d/dropbear
fi
systemctl daemon-reload
service dropbear restart

save_pkg_list post_populate_root_pkgs.txt
touch $GLB_DONE_FILE
sync
#sleep 1

log "[populate-root] EXIT"
#sleep 1

# restart system
#poweroff

